import os, time, glob
from nfstream import NFStreamer
from tqdm import tqdm

def extract(pcap_file_path:str, csv_file_path):
    """Extract features from pcap file.
    """
    streamer = NFStreamer(
        source = pcap_file_path,
        statistical_analysis = True,
        # n_dissections = 0 # comment this line can speed up
    )

    streamer.to_csv(csv_file_path)

def main():
    config = {
        "dataset_path": "DoH_traffic_dataset",
        "csv_dir": "data/output_of_nfstream"
    }

    dataset_path = config["dataset_path"]
    csv_dir = config["csv_dir"]

    domains = os.listdir(dataset_path)

    csv_subdir_path = os.path.join(csv_dir, time.strftime("%Y%m%d%H%M%S", time.localtime()))

    for domain in tqdm(domains, ncols=50):
        csv_domain_dir = os.path.join(csv_subdir_path, domain)
        pcap_domain_dir = os.path.join(dataset_path, domain)
        pcap_files = glob.glob(os.path.join(pcap_domain_dir, "**.pcap"))

        if not os.path.exists(csv_domain_dir):
            os.makedirs(csv_domain_dir)

        for pcap_file in pcap_files:
            cav_file_name = pcap_file.rsplit("/", 1)[-1][:-4] + "csv"
            csv_file_path = os.path.join(csv_subdir_path, domain, cav_file_name)
            extract(pcap_file, csv_file_path)

if __name__ == "__main__":
    main()